# -*- coding: utf-8 -*- 
# Main module:  exploit-mode 

from prompt_toolkit import PromptSession, ANSI
from ropgenerator.exploit.Shellcode import shellcode
from ropgenerator.exploit.Syscall import syscall
from ropgenerator.exploit.Call import call
from ropgenerator.exploit.Pwn import pwn 
from ropgenerator.IO import string_exploit, string_bold, string_special, banner, error

import sys

# Definitions of commands 
CMD_HELP = "help"
CMD_SHELLCODE = "shellcode"
CMD_SYSCALL = "syscall"
CMD_PWN = "pwn"
CMD_CALL = "call"
CMD_MAIN = "main"
CMD_EXIT = "exit"

helpStr = banner([string_bold('Exploit-Mode Commands'),
    string_special('(For more info about a command type <cmd -h>)')])
helpStr += '\n\t' + string_bold(CMD_PWN) + ': \t\tbuild an exploit'
helpStr += '\n\t' + string_bold(CMD_SYSCALL) + ': \tcall a system function'
helpStr += '\n\t' + string_bold(CMD_CALL) + ':\t\tcall a regular function'
helpStr += '\n\t' + string_bold(CMD_SHELLCODE) + ': \tmanage shellcodes'

helpStr += '\n\n\t' + string_bold(CMD_HELP) + ': \t\tshow this help'
helpStr += '\n\t' + string_bold(CMD_MAIN) + ': \t\treturn to the main menu'
helpStr += '\n\t' + string_bold(CMD_EXIT) + ': \t\texit ROPGenerator'

promptSession = PromptSession(ANSI(u"("+ string_exploit(u'exploit') +u")> "))

def exploit_mode():
    """
    Returns
    -------
    True if ROPGenerator must continue
    False if ROPGenerator must be closed 
    """
    
    finish = False
    while( not finish ):
        try:
            # Get command 
            user_input = promptSession.prompt()
            args = user_input.split()
            argslen = len(args)
            if( argslen > 0 ):
                command = args[0]
            else:
                command = None
                continue
                
            # Execute command 
            if( command == CMD_PWN ):
                pwn(args[1:])
            elif( command == CMD_SHELLCODE ):
                shellcode(args[1:])
            elif( command == CMD_SYSCALL ):
                syscall(args[1:])
            elif( command == CMD_CALL ):
                call(args[1:])
            elif( command == CMD_EXIT ):
                return False
            elif( command == CMD_HELP ):
                print(helpStr)
            elif( command == CMD_MAIN ):
                finish = True
            else:
                error("Unknown command '{}'".format(command))
            if( command != None ):
                print('')
        except KeyboardInterrupt:
            pass
    return True












